﻿

#ifndef TENSOR_H
#define TENSOR_H

#include <string>
#include <vector>

template <class DataType = float>
class Tensor
{
public:
    Tensor(const std::vector<int> &dimensions, std::shared_ptr<std::vector<DataType>> buffers, DataType *datas)
    {
        m_dimensions = dimensions;
        m_buffers    = buffers;
        m_datas      = datas;
    }

    Tensor(const std::vector<int> &dimensions)
    {
        m_dimensions = dimensions;
        m_buffers    = std::make_shared<std::vector<DataType>>();
        m_buffers->resize(GetSize());
        m_datas = m_buffers->data();
    }

public:
    inline DataType *GetDatas() { return m_datas; }
    inline int GetSize(int dimension) { return dimension >= 0 && dimension < m_dimensions.size() ? m_dimensions[dimension] : 0; }

public:
    int GetSize()
    {
        if (m_dimensions.size() == 0) return 0;

        int size           = 1;
        int dimension_size = m_dimensions.size();
        for (int i = 0; i < dimension_size; i++)
        {
            size *= m_dimensions[i];
        }
        return size;
    }

protected:
    std::vector<int> m_dimensions;
    std::shared_ptr<std::vector<DataType>> m_buffers;
    DataType *m_datas;
};
#endif